import torch
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import ImageFolder
from model import GoogLeNet, Inception


def test_data():
    root_data = './data/test'
    normalize = transforms.Normalize([0.486, 0.453, 0.415], [0.069, 0.065, 0.067])
    transform = transforms.Compose([transforms.Resize((224, 224)),
                                    transforms.ToTensor(),
                                    normalize])
    data = ImageFolder(root_data, transform=transform)
    test_data = Data.DataLoader(dataset=data,
                                batch_size=1,
                                shuffle=True)
    return test_data


def test(model, test_data):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    test_acc = 0.0
    test_num = 0

    classes = ['猫', '狗']
    count = 0
    # 不计算梯度
    with torch.no_grad():
        for x, y in test_data:
            x = x.to(device)
            y = y.to(device)
            model.eval()
            output = model(x)
            pre_lab = torch.argmax(output, dim=1)
            result = pre_lab.item()
            label = y.item()
            test_acc += torch.sum(pre_lab == y.data)
            test_num += x.size(0)
            if result != label:
                count += 1
                print(f"错误预测{count} : 预测值：{classes[result]} ---------- 真实值：{classes[label]}")

    test_acc = test_acc / test_num * 100
    print("Test Accuracy: {:.2f}%".format(test_acc))


if __name__ == '__main__':
    model = GoogLeNet(Inception)
    # 加载模型参数
    model.load_state_dict(torch.load('checkpoints/classification_epoch_15.pth', map_location='cpu'))
    test_data = test_data()
    test(model, test_data)

